import torch

from ..VideoBenchmark import LVCTestNeural

from PyTorchVideoCompression.mDVC.DVC import DeepVideoCompressor
from PyTorchVideoCompression.DVC.net import VideoCompressor as huDVC
from PyTorchVideoCompression.FVC.net import VideoCompressor as huFVC

from util.metrics import compute_msssim, compute_psnr, mse2psnr

class mDVCTest(LVCTestNeural):
    def __init__(self, seqs, ckpt, intra_quality = 6, gop=12, **kwargs):
        super().__init__(seqs, intra_quality, gop, **kwargs)

        self.inter_model = DeepVideoCompressor()
        self.inter_model.load_state_dict(torch.load(ckpt))
        self.inter_model.eval().cuda()


    def inter_encode(self, x, **kwargs):
        padded_ref = self.dpb["__ref"]
        out = self.inter_model(padded_ref, x)

        return out
    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(output[0])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output[1]["bpp"].item() * self.bpp_ratio

        return {
            "x_hat": output[0],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }


class huDVCTest(LVCTestNeural):
    def __init__(self, seqs, ckpt, intra_quality = 6, gop=12, **kwargs):
        super().__init__(seqs, intra_quality, gop, **kwargs)

        self.inter_model = huDVC()
        self.inter_model.load_state_dict(torch.load(ckpt))
        self.inter_model.eval().cuda()


    def inter_encode(self, x, **kwargs):
        return self.inter_model(x, self.dpb["__ref"])  # ref the latter
        
    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(output[0])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output[7].item() * self.bpp_ratio

        return {
            "x_hat": output[0],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }


class huFVCTest(LVCTestNeural):
    def __init__(self, seqs, ckpt, intra_quality = 6, gop=12, **kwargs):
        super().__init__(seqs, intra_quality, gop, **kwargs)

        self.inter_model = huFVC(2048)
        self.inter_model.load_state_dict(torch.load(ckpt))
        self.inter_model.eval().cuda()


    def inter_encode(self, x):
        return self.inter_model(self.dpb["__ref"], x)  # ref the latter

    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(output[0])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output[1]["bpp"].item() * self.bpp_ratio

        return {
            "x_hat": output[0],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }
